import threading
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs
import json


class VulnHandler(BaseHTTPRequestHandler):
    # 定义漏洞映射表，路径 -> 响应内容
    vuln_responses = {
        "/test/path1": {
            "status": 200,
            "content": "This is a vulnerable response for path1",
            "headers": {"Content-Type": "text/plain"}
        },
        "/test/path2": {
            "status": 200,
            "content": '{"result": "vulnerable", "data": "test"}',
            "headers": {"Content-Type": "application/json"}
        }
    }

    def do_GET(self):
        """处理GET请求"""
        parsed_path = urlparse(self.path)
        query_params = parse_qs(parsed_path.query)

        # 检查路径是否在漏洞映射表中
        if parsed_path.path in self.vuln_responses:
            resp = self.vuln_responses[parsed_path.path]
            self.send_response(resp["status"])
            for header, value in resp["headers"].items():
                self.send_header(header, value)
            self.end_headers()
            self.wfile.write(resp["content"].encode('utf-8'))
        else:
            # self.send_response(404)
            # self.send_header("Content-Type", "text/plain")
            # self.end_headers()
            # self.wfile.write("Not Found".encode('utf-8'))
            return

    def do_POST(self):
        """处理POST请求"""
        content_length = int(self.headers['Content-Length'])
        post_data = self.rfile.read(content_length).decode('utf-8')

        # 检查路径是否在漏洞映射表中
        if self.path in self.vuln_responses:
            resp = self.vuln_responses[self.path]
            self.send_response(resp["status"])
            for header, value in resp["headers"].items():
                self.send_header(header, value)
            self.end_headers()
            self.wfile.write(resp["content"].encode('utf-8'))
        else:
            return
            # self.send_response(404)
            # self.send_header("Content-Type", "text/plain")
            # self.end_headers()
            # self.wfile.write("Not Found".encode('utf-8'))


class VulnMockServer(threading.Thread):
    # 专门用于启动漏洞模拟服务器的类
    def __init__(self, port=1230):
        super().__init__()
        # 服务器端口
        self.__port = port

    def run(self):
        """启动HTTP服务"""
        server_address = ('', self.__port)
        httpd = HTTPServer(server_address, VulnHandler)
        print(f"漏洞模拟服务启动在 http://127.0.0.1:{self.__port}")
        httpd.serve_forever()
